import numpy as np
from typing import List, Tuple, Callable
from scipy.optimize import minimize
from robot_arm import RobotArm

class PathOptimizer:
    def __init__(self, robot: RobotArm):
        self.robot = robot
        self.obstacles = []  # 存储障碍物位置
        
    def add_obstacle(self, obstacle_pos: np.ndarray):
        """添加障碍物位置"""
        self.obstacles.append(obstacle_pos)
        
    def optimize_path(self, 
                     target_position: np.ndarray,
                     initial_angles: List[float] = None,
                     error_weight: float = 1.0,
                     energy_weight: float = 0.0,
                     error_constraint: float = None,
                     collision_avoidance: bool = False,
                     safety_margin: float = 100,
                     num_attempts: int = 10) -> Tuple[List[float], float]:
        """
        优化机器人路径
        
        参数:
            target_position: 目标位置
            initial_angles: 初始关节角度
            error_weight: 误差权重
            energy_weight: 能耗权重
            error_constraint: 误差约束
            collision_avoidance: 是否启用碰撞避免
            safety_margin: 碰撞检测的安全边际(mm)
            num_attempts: 优化尝试次数
            
        返回:
            最优关节角度和目标函数值
        """
        best_result = None
        best_value = float('inf')
        
        # 如果提供了初始角度，先尝试使用它
        if initial_angles is not None:
            result = self._single_optimization(
                target_position,
                initial_angles,
                error_weight,
                energy_weight,
                error_constraint,
                collision_avoidance,
                safety_margin
            )
            if result.success and result.fun < best_value:
                best_result = result
                best_value = result.fun
        
        # 使用随机初始值进行多次优化
        for _ in range(num_attempts):
            initial_angles = self.robot.get_random_angles()
            result = self._single_optimization(
                target_position,
                initial_angles,
                error_weight,
                energy_weight,
                error_constraint,
                collision_avoidance,
                safety_margin
            )
            if result.success and result.fun < best_value:
                best_result = result
                best_value = result.fun
        
        if best_result is None:
            raise ValueError("无法找到满足约束的优化解，请尝试放宽约束或增加尝试次数")
            
        return best_result.x, best_result.fun
    
    def _single_optimization(self,
                           target_position: np.ndarray,
                           initial_angles: List[float],
                           error_weight: float,
                           energy_weight: float,
                           error_constraint: float = None,
                           collision_avoidance: bool = False,
                           safety_margin: float = 100):
        """执行单次优化"""
        bounds = [(np.deg2rad(l), np.deg2rad(u)) for l, u in 
                 zip(self.robot.joint_limits['lower'], 
                     self.robot.joint_limits['upper'])]
        
        # 高碰撞惩罚系数
        collision_penalty = 1e6
        
        def objective(x):
            error = self.robot.calculate_end_effector_error(x, target_position)
            
            # 基础目标函数：误差 + 能耗
            objective_value = error_weight * error
            
            if energy_weight > 0:
                # 计算能耗（如果需要）
                energy = self.calculate_energy(x)
                objective_value += energy_weight * energy
                
            # 如果启用碰撞避免且有障碍物
            if collision_avoidance and self.obstacles:
                # 计算所有关节位置
                joint_positions = self.robot.calculate_joint_positions(x)
                
                # 检查碰撞
                if self.robot.check_collision(joint_positions, self.obstacles, safety_margin):
                    # 如果有碰撞，添加高惩罚值
                    objective_value += collision_penalty
                    
            return objective_value
        
        constraints = []
        
        # 误差约束
        if error_constraint is not None:
            constraints.append({
                'type': 'ineq',
                'fun': lambda x: error_constraint - self.robot.calculate_end_effector_error(x, target_position)
            })
            
        # 碰撞约束（如果启用）
        if collision_avoidance and self.obstacles:
            constraints.append({
                'type': 'ineq',
                'fun': lambda x: 1.0 if not self.robot.check_collision(
                    self.robot.calculate_joint_positions(x), self.obstacles, safety_margin) else -1.0
            })
        
        return minimize(objective, 
                      initial_angles,
                      method='SLSQP',
                      bounds=bounds,
                      constraints=constraints)
    
    def calculate_energy(self, joint_angles: List[float]) -> float:
        """
        计算能耗
        
        根据问题2中的参数计算关节运动的能耗
        """
        # 关节转动惯量
        inertia = [0.5, 0.3, 0.4, 0.6, 0.2, 0.4]
        # 平均角速度
        ang_velocity = [2.0, 1.5, 1.0, 2.5, 3.0, 2.0]
        
        # 计算总能耗 E = Σ(1/2 * I * ω²)
        energy = sum(0.5 * I * (w ** 2) for I, w in zip(inertia, ang_velocity))
        
        return energy 